import numpy as np
import cv2
import sys
from typing import Tuple, Optional


def parse_args() -> Tuple[str, int, int]:
    """Parse and validate command-line arguments."""
    if len(sys.argv) != 4:
        raise ValueError("Usage: python script.py <img_path> <width> <height>")
    
    img_path = sys.argv[1]
    try:
        width = int(sys.argv[2])
        height = int(sys.argv[3])
    except ValueError:
        raise ValueError("Width and height must be integers.")
    
    if width <= 0 or height <= 0:
        raise ValueError("Width and height must be positive integers.")
    
    return img_path, width, height


def load_and_resize_image(img_path: str, width: int, height: int) -> np.ndarray:
    """Load an image and resize it to the specified dimensions."""
    img = cv2.imread(img_path)
    if img is None:
        raise FileNotFoundError(f"Image not found or invalid path: {img_path}")
    
    resized_img = cv2.resize(img, (width, height))
    print(f"Resized image shape: {resized_img.shape}")
    return resized_img


def merge_uv(u: np.ndarray, v: np.ndarray) -> np.ndarray:
    """Interleave U and V channels into a single UV plane (NV12 format)."""
    if u.shape != v.shape:
        raise ValueError("U and V channels must have the same shape.")
    
    # Vectorized interleaving (faster than loops)
    uv = np.empty((u.shape[0], u.shape[1] * 2), dtype=u.dtype)
    uv[:, 0::2] = u  # Even columns = U
    uv[:, 1::2] = v  # Odd columns = V
    return uv

def rgb_to_nv12(image: np.ndarray) -> np.ndarray:
    """Convert BGR image to NV12 (YUV 4:2:0) format."""
    if image.ndim != 3 or image.shape[2] != 3:
        raise ValueError("Input must be a 3-channel BGR image.")
    # Convert BGR to YUV (OpenCV uses BGR order)
    yuv = cv2.cvtColor(image, cv2.COLOR_BGR2YUV_I420)
    h, w = image.shape[:2]
    
    # 计算各分量的大小
    y_size = h * w
    uv_size = y_size // 4

    # 分解 YUV 分量
    yuv = yuv.reshape(-1)
    y = yuv[:y_size].reshape(h, w)
    u = yuv[y_size:y_size+uv_size].reshape(h//2, w//2)
    v = yuv[y_size+uv_size:y_size+2*uv_size].reshape(h//2, w//2)

    # For NV12, we interleave U and V planes
    uv = np.empty((h//2, w), dtype=np.uint8)
    uv[:, 0::2] = u
    uv[:, 1::2] = v
    
    # Combine Y and interleaved UV planes
    nv12 = np.vstack((y, uv))
    
    return nv12

def rgb_to_I420(image: np.ndarray) -> np.ndarray:
    """Convert BGR image to NV12 (YUV 4:2:0) format."""
    if image.ndim != 3 or image.shape[2] != 3:
        raise ValueError("Input must be a 3-channel BGR image.")
    # Convert BGR to YUV (OpenCV uses BGR order)
    yuv = cv2.cvtColor(image, cv2.COLOR_BGR2YUV_I420)
    h, w = image.shape[:2]
    
    
    # I420 (YUV420P) to NV12 (YUV420SP) conversion
    y = yuv[:h, :]
    uv = yuv[h:, :]
    
    # Reshape UV plane (I420: U and V are separate, NV12: interleaved)
    uv = uv.reshape((h // 2, w))
    return np.concatenate([y, uv])


def save_nv12(data: np.ndarray, original_shape: Tuple[int, int]) -> None:
    """Save NV12 data to a binary file."""
    output_path = f"test_w{original_shape[1]}_h{original_shape[0]}.yuv"
    data.tofile(output_path)
    print(f"NV12 data saved to: {output_path}")

def save_jpg(data: np.ndarray, original_shape: Tuple[int, int]) -> None:
    output_path = f"test_w{original_shape[1]}_h{original_shape[0]}.jpg"
    cv2.imwrite(output_path,data)
    print(f"JPG saved to: {output_path}")

def main() -> None:
    try:
        # Step 1: Parse arguments
        img_path, width, height = parse_args()
        
        # Step 2: Load and resize image
        img = load_and_resize_image(img_path, width, height)
        save_jpg(img, img.shape[:2])
        
        # Step 3: Convert to NV12
        # nv12_data = rgb_to_I420(img)
        nv12_data = rgb_to_nv12(img)

        # Step 4: Save output
        save_nv12(nv12_data, img.shape[:2])
    
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()
    
# python jpg-to-nv12.py test.jpg 640 640
